#include <Windows.h>
#include <winsock2.h>
#include <pcap.h>
#include <stdio.h>
#include <time.h>
#include <string.h>


#pragma comment(lib, "wpcap.lib")
#pragma comment(lib, "Ws2_32.lib")

#define CSV_FILE "traffic_log.csv"

// 统计结构体
typedef struct {
    char mac[18];
    char ip[16];
    long total_bytes;
} StatEntry;

StatEntry src_stats[1000]; 
int stat_count = 0;

// 写入CSV日志
void write_csv(const char* time_str, const char* src_mac, const char* src_ip,
    const char* dst_mac, const char* dst_ip, int frame_len) {
    FILE* fp = fopen(CSV_FILE, "a");
    if (fp) {
        fprintf(fp, "%s,%s,%s,%s,%s,%d\n",
            time_str, src_mac, src_ip, dst_mac, dst_ip, frame_len);
        fclose(fp);
    }
}

// 更新统计
void update_stat(const char* mac, const char* ip, int len) {
    for (int i = 0; i < stat_count; i++) {
        if (strcmp(src_stats[i].mac, mac) == 0 && strcmp(src_stats[i].ip, ip) == 0) {
            src_stats[i].total_bytes += len;
            return;
        }
    }
    // 新增条目
    if (stat_count < 1000) {
        strcpy_s(src_stats[stat_count].mac, sizeof(src_stats[stat_count].mac), mac);
        strcpy_s(src_stats[stat_count].ip, sizeof(src_stats[stat_count].ip), ip);
        src_stats[stat_count].total_bytes = len;
        stat_count++;
    }
}

// 定时输出统计（需多线程）
DWORD WINAPI timer_thread(LPVOID lpParam) {
    while (1) {
        Sleep(60000); // 60秒
        printf("\n--- Traffic Statistics (Last 1 Min) ---\n");
        for (int i = 0; i < stat_count; i++) {
            printf("MAC: %-18s | IP: %-15s | Sent: %ld bytes\n",
                src_stats[i].mac, src_stats[i].ip, src_stats[i].total_bytes);
        }
        // 清空统计
        memset(src_stats, 0, sizeof(src_stats));
        stat_count = 0;
    }
    return 0;
}

// 抓包回调函数
void packet_handler(u_char* param, const struct pcap_pkthdr* header, const u_char* pkt_data) {
    struct tm ltime;
    char timestr[20];
    time_t local_tv_sec = header->ts.tv_sec;
    localtime_s(&ltime, &local_tv_sec);
    strftime(timestr, sizeof(timestr), "%Y-%m-%d %H:%M:%S", &ltime);

    // 解析以太网帧
    struct ethhdr {
        u_char h_dest[6];
        u_char h_source[6];
        u_short h_proto;
    } *eth = (struct ethhdr*)pkt_data;

    char src_mac[18], dst_mac[18];
    sprintf_s(src_mac, sizeof(src_mac), "%02X:%02X:%02X:%02X:%02X:%02X",
        eth->h_source[0], eth->h_source[1], eth->h_source[2],
        eth->h_source[3], eth->h_source[4], eth->h_source[5]);
    sprintf_s(dst_mac, sizeof(dst_mac), "%02X:%02X:%02X:%02X:%02X:%02X",
        eth->h_dest[0], eth->h_dest[1], eth->h_dest[2],
        eth->h_dest[3], eth->h_dest[4], eth->h_dest[5]);

    // 仅处理IPv4
    if (ntohs(eth->h_proto) == 0x0800) {
        struct iphdr {
            unsigned char ihl : 4;
            unsigned char version : 4;
            u_char tos;
            u_short tot_len;
            u_short id;
            u_short frag_off;
            u_char ttl;
            u_char protocol;
            u_short check;
            u_int saddr;
            u_int daddr;
        } *ip = (struct iphdr*)(pkt_data + 14); // 以太网帧头14字节

        char src_ip[16], dst_ip[16];
        inet_ntop(AF_INET, &ip->saddr, src_ip, sizeof(src_ip));
        inet_ntop(AF_INET, &ip->daddr, dst_ip, sizeof(dst_ip));

        // 写入日志
        write_csv(timestr, src_mac, src_ip, dst_mac, dst_ip, header->len);

        // 更新统计
        update_stat(src_mac, src_ip, header->len);
    }
}

int main() {
    pcap_if_t* alldevs;
    char errbuf[PCAP_ERRBUF_SIZE];

    // 获取网卡列表
    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        printf("Error finding devices: %s\n", errbuf);
        return 1;
    }

    // 选择第一个网卡
    pcap_t* handle = pcap_open_live(alldevs->name, 65535, 1, 1000, errbuf);
    if (!handle) {
        printf("Error opening device: %s\n", errbuf);
        return 1;
    }

    // 启动统计线程
    CreateThread(NULL, 0, timer_thread, NULL, 0, NULL);

    // 开始抓包
    pcap_loop(handle, 0, packet_handler, NULL);

    pcap_close(handle);
    pcap_freealldevs(alldevs);
    return 0;
}